#include <cstdio>
#include <systemc>
#include <tlm.h>
#include <tlm_utils/simple_initiator_socket.h>
#include <tlm_utils/simple_target_socket.h>
#include <tlm_utils/tlm_quantumkeeper.h>

using namespace sc_core;
using namespace sc_dt;
using namespace std;

// User-defined memory manager, which maintains a pool of transactions
class MemoryManager : public tlm::tlm_mm_interface {
    typedef tlm::tlm_generic_payload gp_t;
public:
    MemoryManager() {}
    ~MemoryManager() {
        for (std::size_t i = 0; i < all_trans.size(); ++i) {
            delete all_trans[i];
        }
    }
    // Allocate a new transaction or reuse one from the free pool
    gp_t* allocate() {
        gp_t* ptr;
        if (!free_pool.empty()) {
            ptr = free_pool.back();
            free_pool.pop_back();
        } else {
            ptr = new gp_t(this);
            all_trans.push_back(ptr);
        }
        // printf("------ Called MemoryManager.allocate(), trans_count = %zu\n", ++count);
        return ptr;
    }
    // Free a transaction back to the memory manager
    void free(gp_t* trans) {
        trans->reset();
        free_pool.push_back(trans);
        // printf("------ Called MemoryManager.free(), trans_count = %zu\n", --count);
    }
private:
    std::vector<gp_t*> free_pool;
    std::vector<gp_t*> all_trans;
    size_t count = 0;
};

const int RUN_LENGTH = 256;

// Master1 writes to all 4 memories, and demonstrates DMI and debug transport
// Does not use an explicit memory manager
SC_MODULE(Master1) {
    tlm_utils::simple_initiator_socket<Master1> socket;

    SC_CTOR(Master1) : socket("socket"), dmi_ptr_valid(false) {
        socket.register_invalidate_direct_mem_ptr(this, &Master1::invalidate_direct_mem_ptr);

        SC_THREAD(thread_process);

        // All initiators use a quantum of 1us, that is, they synchronize
        // themselves to simulation time every 1us using the quantum keeper
        qk.set_global_quantum(sc_time(1, SC_US));
        qk.reset();
    }

    void thread_process() {
        // Use debug transaction interface to dump entire memory contents
        dump();

        tlm::tlm_generic_payload *trans = new tlm::tlm_generic_payload;
        sc_time delay, delta;
        bool use_dmi = false;

        for (int i = 0; i < RUN_LENGTH; i += 4) {
            data = i;
            delay = qk.get_local_time();
            use_dmi = false;

            // Use DMI
            if (dmi_ptr_valid &&
                sc_dt::uint64(i) >= dmi_data.get_start_address() &&
                sc_dt::uint64(i) <= dmi_data.get_end_address()) {
                use_dmi = true;
                // Bypass transport interface and use direct memory interface
                assert(dmi_data.is_write_allowed());
                memcpy(dmi_data.get_dmi_ptr() + i - dmi_data.get_start_address(), &data, 4);

                // Accumulate memory latency into local time
                delta = dmi_data.get_read_latency();
                delay += delta;
            }
            // No DMI, so use blocking transport interface
            else {
                trans->set_command(tlm::TLM_WRITE_COMMAND);
                trans->set_address(i);
                trans->set_data_ptr(reinterpret_cast<unsigned char *>(&data));
                trans->set_data_length(4);
                trans->set_streaming_width(4); // = data_length to indicate no streaming
                trans->set_byte_enable_ptr(0); // 0 indicates unused
                trans->set_dmi_allowed(false); // Mandatory initial value
                trans->set_response_status(tlm::TLM_INCOMPLETE_RESPONSE); // Mandatory initial value

                sc_time delay_prev = delay;
                socket->b_transport(*trans, delay);
                delta = delay - delay_prev;

                // Initiator obliged to check response status
                if (trans->is_response_error()) {
                    SC_REPORT_ERROR("TLM-2", trans->get_response_string().c_str());
                }
                if (trans->is_dmi_allowed()) {
                    dmi_data.init();
                    dmi_ptr_valid = socket->get_direct_mem_ptr(*trans, dmi_data);
                }
            }

            // Accumulate local time and synchronize when quantum is reached
            qk.set(delay);
            printf("WRITE%s addr = %02x, data = %02x, global_time=%s, local_time=%s, delta_delay=%s\n",
                (use_dmi ? "/DMI" : "    "), i, data, sc_time_stamp().to_string().c_str(), delay.to_string().c_str(), delta.to_string().c_str());

            qk.inc(sc_time(100, SC_NS)); // Model time used for additional processing
            if (qk.need_sync()) {
                printf("==> Master1 Synchronized, global_time=%s, local_time=%s, compute_local_quantum=%s\n",
                    sc_time_stamp().to_string().c_str(), qk.get_local_time().to_string().c_str(),
                    tlm::tlm_global_quantum::instance().compute_local_quantum().to_string().c_str());
                qk.sync();
            }
        }

        // Use debug transaction interface to dump entire memory contents
        dump();
    }

    virtual void invalidate_direct_mem_ptr(sc_dt::uint64 start_range,
                                           sc_dt::uint64 end_range) {
        printf("INVALIDATE DMI [%02x, %02x] for Master1 at %s\n",
               (uint32_t)start_range, (uint32_t)end_range, sc_time_stamp().to_string().c_str());
        // Ignore range and invalidate all DMI pointers regardless
        dmi_ptr_valid = false;
    }

    void dump() {
        uint8_t buffer[64];

        // Use debug transaction interface to dump memory contents
        printf("Dump memory contents using debug transport at time %s\n", sc_time_stamp().to_string().c_str());

        for (uint32_t k = 0; k < 4; k++) {
            tlm::tlm_generic_payload dbg;
            sc_dt::uint64 A = 64 * k;
            dbg.set_address(A);
            dbg.set_read();
            dbg.set_data_length(64);
            dbg.set_data_ptr(buffer);

            uint32_t n_bytes = socket->transport_dbg(dbg);

            for (uint32_t i = 0; i < n_bytes; i += 4) {
                printf("mem[%02x] = %08x\n", (uint32_t)(A + i),
                       *(reinterpret_cast<uint32_t *>(&buffer[i])));
            }
        }
        printf("\n");
    }

    int data; // Internal data buffer used by initiator with generic payload
    tlm_utils::tlm_quantumkeeper qk; // Quantum keeper for temporal decoupling
    bool dmi_ptr_valid;
    tlm::tlm_dmi dmi_data;
};


// Master2 reads from all 4 memories, but does not use DMI or debug transport
// Uses an explicit memory manager and transaction pool
SC_MODULE(Master2) {
    tlm_utils::simple_initiator_socket<Master2> socket;

    SC_CTOR(Master2) : socket("socket") {
        // No callback methods registered with socket
        SC_THREAD(thread_process);
        qk.set_global_quantum(sc_time(1, SC_US));
    }

    void thread_process() {
        tlm::tlm_generic_payload *trans;
        sc_time delay, delta;

        // Make sure Master2 starts after Master1
        wait(1, SC_US);
        // Reset the local quantum keeper !!! MUST
        qk.reset();

        for (int i = 0; i < RUN_LENGTH; i += 4) {
            // Grab a new transaction from the memory manager
            trans = mm.allocate();
            trans->acquire();

            data = i;

            trans->set_command(tlm::TLM_READ_COMMAND);
            trans->set_address(i);
            trans->set_data_ptr(reinterpret_cast<unsigned char *>(&data));
            trans->set_data_length(4);
            trans->set_streaming_width(4); // = data_length to indicate no streaming
            trans->set_byte_enable_ptr(0); // 0 indicates unused
            trans->set_dmi_allowed(false); // Mandatory initial value
            trans->set_response_status(tlm::TLM_INCOMPLETE_RESPONSE); // Mandatory initial value

            delay = qk.get_local_time();
            sc_time delay_prev = delay;
            socket->b_transport(*trans, delay);
            delta = delay - delay_prev;

            // Initiator obliged to check response status
            if (trans->is_response_error()) {
                SC_REPORT_ERROR("TLM-2", trans->get_response_string().c_str());
            }
            if (data != i) {
                SC_REPORT_ERROR("TLM-2", "Mismatch in initiator when reading back data");
            }

            trans->release();

            // Accumulate local time and synchronize when quantum is reached
            qk.set(delay);
            printf("READ      addr = %02x, data = %02x, global_time=%s, local_time=%s, delta_delay=%s\n",
                i, data, sc_time_stamp().to_string().c_str(), delay.to_string().c_str(), delta.to_string().c_str());

            qk.inc(sc_time(100, SC_NS)); // Model time used for additional processing
            if (qk.need_sync()) {
                printf("==> Master2 Synchronized, global_time=%s, local_time=%s, compute_local_quantum=%s\n",
                    sc_time_stamp().to_string().c_str(), qk.get_local_time().to_string().c_str(),
                    tlm::tlm_global_quantum::instance().compute_local_quantum().to_string().c_str());
                qk.sync();
            }
        }
    }

    int data; // Internal data buffer used by initiator with generic payload

    MemoryManager mm;                // Memory manager
    tlm_utils::tlm_quantumkeeper qk; // Quantum keeper for temporal decoupling
};

// Bus model supports multiple initiators and multiple targets
template <unsigned int N_INITIATORS, unsigned int N_TARGETS>
SC_MODULE(Bus) {
    // Tagged sockets allow incoming transactions to be identified
    // Connect to Intiators
    tlm_utils::simple_target_socket_tagged<Bus> *target_socket[N_INITIATORS];
    // Connect to Targets
    tlm_utils::simple_initiator_socket_tagged<Bus> *initiator_socket[N_TARGETS];

    SC_CTOR(Bus) {
        char name[128];
        for (int i = 0; i < N_INITIATORS; i++) {
            sprintf(name, "target_socket_%d", i);
            target_socket[i] = new tlm_utils::simple_target_socket_tagged<Bus>(name);
            target_socket[i]->register_b_transport(this, &Bus::b_transport, i);
            target_socket[i]->register_get_direct_mem_ptr(this, &Bus::get_direct_mem_ptr, i);
            target_socket[i]->register_transport_dbg(this, &Bus::transport_dbg, i);
        }
        for (int i = 0; i < N_TARGETS; i++) {
            sprintf(name, "initiator_socket_%d", i);
            initiator_socket[i] = new tlm_utils::simple_initiator_socket_tagged<Bus>(name);
            initiator_socket[i]->register_invalidate_direct_mem_ptr(this, &Bus::invalidate_direct_mem_ptr, i);
        }
    }

    // Tagged TLM-2.0 blocking transport method [Forward Path]
    virtual void b_transport(int id, tlm::tlm_generic_payload &trans, sc_time &delay) {
        if (id >= N_INITIATORS) {
            SC_REPORT_FATAL("TLM-2", "Invalid tagged socket id in bus");
            return;
        }
        sc_dt::uint64 address = trans.get_address();
        sc_dt::uint64 masked_address;
        unsigned int target_nr = decode_address(address, masked_address);

        if (target_nr < N_TARGETS) {
            // Modify address within transaction
            trans.set_address(masked_address);

            // Forward transaction to appropriate target
            (*initiator_socket[target_nr])->b_transport(trans, delay);

            // Replace original address
            trans.set_address(address);
        }
    }

    // Tagged TLM-2.0 forward DMI method
    virtual bool get_direct_mem_ptr(int id, tlm::tlm_generic_payload &trans, tlm::tlm_dmi &dmi_data) {
        sc_dt::uint64 masked_address;
        unsigned int target_nr = decode_address(trans.get_address(), masked_address);
        if (target_nr >= N_TARGETS)  return false;

        trans.set_address(masked_address);

        bool status = (*initiator_socket[target_nr])->get_direct_mem_ptr(trans, dmi_data);

        // Calculate DMI address of target in system address space
        dmi_data.set_start_address(compose_address(target_nr, dmi_data.get_start_address()));
        dmi_data.set_end_address(compose_address(target_nr, dmi_data.get_end_address()));

        return status;
    }

    // Tagged debug transaction method
    virtual unsigned int transport_dbg(int id, tlm::tlm_generic_payload &trans) {
        sc_dt::uint64 masked_address;
        unsigned int target_nr = decode_address(trans.get_address(), masked_address);
        if (target_nr >= N_TARGETS) return 0;
        trans.set_address(masked_address);

        // Forward debug transaction to appropriate target
        return (*initiator_socket[target_nr])->transport_dbg(trans);
    }

    // Tagged backward DMI method
    virtual void invalidate_direct_mem_ptr(int id, sc_dt::uint64 start_range, sc_dt::uint64 end_range) {
        // Reconstruct address range in system memory map
        sc_dt::uint64 bw_start_range = compose_address(id, start_range);
        sc_dt::uint64 bw_end_range = compose_address(id, end_range);

        // Propagate call backward to all initiators
        for (unsigned int i = 0; i < N_INITIATORS; i++)
            (*target_socket[i])->invalidate_direct_mem_ptr(bw_start_range, bw_end_range);
    }

    // Simple fixed address decoding
    inline unsigned int decode_address(sc_dt::uint64 address, sc_dt::uint64 &masked_address) {
        unsigned int target_nr = static_cast<unsigned int>((address >> 6) & 0x3);
        masked_address = address & 0x3F;
        return target_nr;
    }

    inline sc_dt::uint64 compose_address(unsigned int target_nr, sc_dt::uint64 address) {
        return (target_nr << 6) | (address & 0x3F);
    }
};

// Target memory implements b_transport, DMI and debug
SC_MODULE(Memory) {
    tlm_utils::simple_target_socket<Memory> socket;

    enum { SIZE = 64 };
    const sc_time LATENCY;

    SC_CTOR(Memory) : socket("socket"), LATENCY(10, SC_NS) {
        socket.register_b_transport(this, &Memory::b_transport);
        socket.register_get_direct_mem_ptr(this, &Memory::get_direct_mem_ptr);
        socket.register_transport_dbg(this, &Memory::transport_dbg);

        // Initialize memory with random data
        for (int i = 0; i < SIZE; i++) {
            mem[i] = 0xAA000000 | (mem_nr << 20) | (rand() % 256);
        }

        // Each instance is given identifiable contents to help debug
        ++mem_nr;

        SC_THREAD(invalidation_process);
    }

    virtual void b_transport(tlm::tlm_generic_payload &trans, sc_time &delay) {
        tlm::tlm_command cmd = trans.get_command();
        sc_dt::uint64 addr = trans.get_address() / 4;
        unsigned char *ptr = trans.get_data_ptr();
        unsigned int len = trans.get_data_length();
        unsigned char *byte = trans.get_byte_enable_ptr();
        unsigned int width = trans.get_streaming_width();

        if (addr > sc_dt::uint64(SIZE)) {
            trans.set_response_status(tlm::TLM_ADDRESS_ERROR_RESPONSE);
            return;
        }
        if (byte != 0) {
            trans.set_response_status(tlm::TLM_BYTE_ENABLE_ERROR_RESPONSE);
            return;
        }
        if (len > 4 || width < len) {
            trans.set_response_status(tlm::TLM_BURST_ERROR_RESPONSE);
            return;
        }

        if (trans.get_command() == tlm::TLM_READ_COMMAND)
            memcpy(ptr, &mem[addr], len);
        else if (cmd == tlm::TLM_WRITE_COMMAND)
            memcpy(&mem[addr], ptr, len);

        // Use temporal decoupling: add memory latency to delay argument
        delay += LATENCY;

        trans.set_dmi_allowed(true);
        trans.set_response_status(tlm::TLM_OK_RESPONSE);
    }

    // TLM-2.0 DMI method
    virtual bool get_direct_mem_ptr(tlm::tlm_generic_payload &trans, tlm::tlm_dmi &dmi_data) {
        // Permit read and write access
        dmi_data.allow_read_write();

        // Set other details of DMI region
        dmi_data.set_dmi_ptr(reinterpret_cast<unsigned char *>(&mem[0]));
        dmi_data.set_start_address(0);
        dmi_data.set_end_address(SIZE * 4 - 1);
        dmi_data.set_read_latency(LATENCY);
        dmi_data.set_write_latency(LATENCY);

        return true;
    }

    // TLM-2.0 debug transaction method
    virtual unsigned int transport_dbg(tlm::tlm_generic_payload &trans) {
        tlm::tlm_command cmd = trans.get_command();
        sc_dt::uint64 addr = trans.get_address() / 4;
        unsigned char *ptr = trans.get_data_ptr();
        unsigned int len = trans.get_data_length();

        // Calculate the number of bytes to be actually copied
        unsigned int num_bytes = (len < (SIZE - addr) * 4) ? len : (SIZE - addr) * 4;

        if (cmd == tlm::TLM_READ_COMMAND) {
            memcpy(ptr, &mem[addr], num_bytes);
        }
        else if (cmd == tlm::TLM_WRITE_COMMAND) {
            memcpy(&mem[addr], ptr, num_bytes);
        }

        return num_bytes;
    }

    void invalidation_process() {
        // Invalidate DMI pointers once just as an example of routing a call back to initiators
        wait(3, SC_US);
        socket->invalidate_direct_mem_ptr(0, SIZE - 1);
    }

    int mem[SIZE];
    static unsigned int mem_nr; // Unique memory number to help debug
};

unsigned int Memory::mem_nr = 0;
SC_MODULE(Top) {
    Master1 *master1;
    Master2 *master2;
    Bus<2, 4> *bus;
    Memory *memory[4];

    SC_CTOR(Top) {
        master1 = new Master1("master1");
        master2 = new Master2("master2");
        bus = new Bus<2, 4>("bus");

        master1->socket.bind(*(bus->target_socket[0]));
        master2->socket.bind(*(bus->target_socket[1]));

        char name[128];
        for (int i = 0; i < 4; i++) {
            sprintf(name, "memory_%d", i);
            memory[i] = new Memory(name);

            (*(bus->initiator_socket[i])).bind(memory[i]->socket);
        }
    }
};

int sc_main(int argc, char *argv[]) {
    Top top("top");
    sc_start();
    return 0;
}
